import os
import time
import datetime
from PIL import ImageGrab
import threading
import signal
import sys

class AutoScreenshot:
    def __init__(self, interval=2, save_dir="ss"):
        """
        初始化自动截屏工具
        
        Args:
            interval: 截屏间隔时间（秒）
            save_dir: 保存截图的文件夹名称
        """
        self.interval = interval
        self.save_dir = save_dir
        self.running = False
        self.thread = None
        
        # 创建保存截图的文件夹
        if not os.path.exists(self.save_dir):
            os.makedirs(self.save_dir)
            print(f"已创建截图保存文件夹: {self.save_dir}")
    
    def take_screenshot(self):
        """截取屏幕并保存"""
        try:
            # 获取当前时间作为文件名
            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
            filename = f"screenshot_{timestamp}.png"
            filepath = os.path.join(self.save_dir, filename)
            
            # 截取整个屏幕
            screenshot = ImageGrab.grab()
            
            # 保存截图
            screenshot.save(filepath)
            print(f"已保存截图: {filename}")
            
        except Exception as e:
            print(f"截图失败: {str(e)}")
    
    def start(self):
        """开始自动截屏"""
        if self.running:
            print("自动截屏已在运行中")
            return
            
        self.running = True
        print("开始自动截屏...")
        
        # 创建并启动线程
        self.thread = threading.Thread(target=self._run_screenshot_loop)
        self.thread.daemon = True
        self.thread.start()
    
    def stop(self):
        """停止自动截屏"""
        if not self.running:
            return
            
        self.running = False
        print("正在停止自动截屏...")
        
        # 等待线程结束
        if self.thread and self.thread.is_alive():
            self.thread.join(timeout=2)
        
        print("自动截屏已停止")
        sys.exit(0)
    
    def _run_screenshot_loop(self):
        """运行截屏循环"""
        while self.running:
            self.take_screenshot()
            
            # 等待下一个截屏间隔
            start_time = time.time()
            while self.running and (time.time() - start_time) < self.interval:
                time.sleep(0.1)

def signal_handler(signum, frame):
    """处理Ctrl+C信号"""
    print("\n检测到Ctrl+C，正在退出程序...")
    auto_ss.stop()

def main():
    """主函数 - 启动后立即开始自动截屏"""
    print("=== 自动截屏工具 ===")
    print("功能：每2秒自动截取整个屏幕")
    print("保存路径：当前目录下的 screenshot 文件夹")
    print("操作说明：按 Ctrl+C 停止程序")
    print("-" * 30)
    
    # 创建自动截屏实例
    auto_ss = AutoScreenshot(interval=2, save_dir="screenshot")
    
    # 设置信号处理
    signal.signal(signal.SIGINT, signal_handler)
    
    # 启动自动截屏
    auto_ss.start()
    
    try:
        # 保持程序运行
        while auto_ss.running:
            time.sleep(1)
    except KeyboardInterrupt:
        signal_handler(None, None)

if __name__ == "__main__":
    main()